import cv2
import numpy as np

# 在HSV色彩空间下得到二值图
def Get_HSV(image,hmin,hmax,smin,smax,vmin,vmax,color):

    # 1 to HSV
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    h, s, v = cv2.split(hsv)

    v = cv2.equalizeHist(v)

    # 3 set threshold (binary image)
    # if value in (min, max):white; otherwise:black
    h_binary = cv2.inRange(np.array(h), np.array(hmin), np.array(hmax))
    s_binary = cv2.inRange(np.array(s), np.array(smin), np.array(smax))
    v_binary = cv2.inRange(np.array(v), np.array(vmin), np.array(vmax))

    # 4 get binary（对H、S、V三个通道分别与操作）
    binary = cv2.bitwise_and(h_binary, cv2.bitwise_and(s_binary, v_binary))

    return binary


#图形经行基本处理
def image_process(binary):

    blur = cv2.GaussianBlur(binary, (9, 9), 0)

    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (13,13))
    Open = cv2.morphologyEx(blur, cv2.MORPH_OPEN, kernel)

    Close = cv2.morphologyEx(Open, cv2.MORPH_CLOSE, kernel)
    Close = cv2.morphologyEx(Close, cv2.MORPH_CLOSE, kernel)
#    edged = cv2.Canny(Close,100,200)
    
    ret,thresh = cv2.threshold(Close,200,255,cv2.THRESH_TOZERO)

    return thresh


#对于图形轮廓检测
def findcontours_circle(edged,img):

    x = 0
    y = 0
    w = 0
    h = 0
    contours, hierarchy = cv2.findContours(edged, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    #只找出最大的轮廓
    contours = sorted(contours,key=cv2.contourArea,reverse = True)

    for contour in contours:
        #计算轮廓所围成面积
        area = cv2.contourArea(contour)
        if area < 100:
            continue

        #判断轮廓线是否为闭合图形
        perimeter = cv2.arcLength(contour,True)

        #轮廓逼近
        approx = cv2.approxPolyDP(contour,0.02*perimeter,True)

        #边数
        vertices = len(approx)

        if vertices > 7:

            screenCnt = approx

            #画出图形
#            cv2.drawContours(img,[screenCnt],-1,(0,0,255),2)       
            x,y,w,h = cv2.boundingRect(screenCnt)

            cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
        
    return img,x,y,w,h




#寻找不同颜色小球小球
#第一个参数用于计算，第二个参数用来画图
def Find_blud_ball(image,img):

    binary_blue = Get_HSV(image,80,114,134,255,0,255,"blue")

 #   cv2.imshow("blue", binary_blue) 

    edged_blue = image_process(binary_blue)

 #   cv2.imshow("edged_blue", edged_blue) 

    img,x,y,w,h = findcontours_circle(edged_blue,img)

    return img,x,y,w,h

def Find_red_ball(image,img):

    binary_red = Get_HSV(image,156,179,43,255,46,255,"red")

#    cv2.imshow("red", binary_red) 

    edged_red = image_process(binary_red)

#    cv2.imshow("edged_red", edged_red) 

    img,x,y,w,h = findcontours_circle(edged_red,img)

    return img,x,y,w,h

def Find_yellow_ball(image,img):

    binary_yellow = Get_HSV(image,22,33,43,255,0,255,"yellow")

#    cv2.imshow("yellow", binary_yellow) 

    edged_yellow = image_process(binary_yellow)

#    cv2.imshow("edged_yellow", edged_yellow) 

    img,x,y,w,h = findcontours_circle(edged_yellow,img)

    return img,x,y,w,h

def Find_black_ball(image,img):

    binary_black = Get_HSV(image,0,179,0,255,0,50,"black")

#    cv2.imshow("black", binary_black) 

    edged_black = image_process(binary_black)

    img,x,y,w,h = findcontours_circle(edged_black,img)

    return img,x,y,w,h


#观察矿仓中是否有球
#第一个参数用于计算，第二个参数用来画图
def Find_box_ball(image,img,color):

    if(color == 0):
        img,x1,y1,w1,h1 = Find_blud_ball(image,img)
        img,x2,y2,w2,h2 = Find_yellow_ball(image,img)
    else:
        img,x1,y1,w1,h1 = Find_red_ball(image,img) 
        img,x2,y2,w2,h2 = Find_yellow_ball(image,img)

    if(w1*h1 > 100 or w2*h2 > 100):
        return 1
    else:
        return 0



#0找蓝球 1找红球
color = 0
#矿仓是否有球
box_state = 0
#
z=0
speed=0

capture = cv2.VideoCapture(1)  # 打开笔记本内置摄像头
while (capture.isOpened()):  # 笔记本内置摄像头被打开后
    retval, image = capture.read(1)  # 从摄像头中实时读取视频
    #img用做绘图，image用作处理
    img = image.copy()



    #识别矿仓
    box = image[300:450,50:600]
    #box_draw用做绘图，box用做计算
    box_draw =box.copy()
    box_state = Find_box_ball(box,box_draw,color)

    cv2.imshow("box", box_draw)



    #识别小球
    if(color == 0):
        img,x1,y1,w1,h1 = Find_blud_ball(image,img)
        area1 = w1*h1 
        img,x2,y2,w2,h2 = Find_yellow_ball(image,img)
        area2 = w2*h2
    else:    
        img,x1,y1,w1,h1 = Find_red_ball(image,img)


    if(area1 != 0 or area2 != 0):
        if(area1 > area2):
            z = x1 - 240 
        else:
            z = x2 - 240  

        if(box_state == 0):
           speed=10 
        else:
            speed=5
    else:
        z = 10




    cv2.imshow("Video", img)  # 在窗口中显示读取到的视频
    
    key = cv2.waitKey(1)  # 窗口的图像刷新时间为1毫秒
    if key == 27:  # 如果按下esc
        break
capture.release()  # 关闭笔记本内置摄像头
cv2.destroyAllWindows()  # 销毁显示摄像头视频的窗口


